#include "acl/acl.h"
#include "../../include/custom_type.h"
#include "../data_utils.h"
#include "../run_main.h"

#include "../LLMsGEMM_batch_QKT/LLMsGEMM_batch_QKT_host.h"
#include "../LLMsGEMM_batch_QKTV/LLMsGEMM_batch_QKTV_host.h"
#include "../LLMsGEMM_batch_QKTVP/LLMsGEMM_batch_QKTVP_host.h"

#include <cstdio>
#include <string>

int main(int argc, char ** argv){
    int64_t wholeM = std::stoi( argv[1] );
    int64_t wholeN = std::stoi( argv[2] );
    int64_t wholeK = std::stoi( argv[3] );
    int64_t batchCount = std::stoi( argv[4] );

    // 打印输出 layoutA layoutB zeroPaddingM zeroPaddingN zeroPaddingK batchCount
    printf("\n\n"
        "wholeM: %ld, wholeN: %ld, wholeK: %ld, " 
        "batchCount: %ld. \n\n", 
        wholeM, wholeN, wholeK, 
        batchCount);


    // acl初始化
    int device = 4;
    const char *aclConfigPath = "../acl.json";
    ACL_CHECK(aclInit(/*aclConfigPath*/nullptr));
    ACL_CHECK(aclrtSetDevice(device));
    aclrtStream stream;
    ACL_CHECK(aclrtCreateStream(&stream));

    std::string src="./data/";

    std::string path_alpha = src + "alpha.bin";
    std::string path_beta  = src + "beta.bin";
    std::string path_validM = src + "validM.bin";

    // =========================================================================

    std::string path_Q  = src + "Q.bin";
    std::string path_KT = src + "KT.bin";
    std::string path_C1 = src + "C1.bin";
    std::string path_expect_result1 = src + "expect_result1.bin";

    half* pre_output1 = nullptr;
    half* cur_output1 = nullptr;

    run_main(
        LLMsGEMM_batch_QKT_host,
        stream, 
        RowMajor, 
        RowMajor, 
        wholeM, 
        wholeM, 
        wholeK, 
        batchCount, 
        path_Q, 
        path_KT, 
        path_C1, 
        path_alpha, 
        path_beta, 
        path_validM, 
        path_expect_result1, 
        pre_output1, 
        cur_output1, 
        1, 
        0
    );

    // =========================================================================

    std::string path_A2 = "";
    std::string path_V  = src + "V.bin";
    std::string path_C2 = src + "C2.bin";
    std::string path_expect_result2 = src + "expect_result2.bin";

    half* pre_output2 = cur_output1;
    half* cur_output2 = nullptr;

    run_main(
        LLMsGEMM_batch_QKTV_host,
        stream, 
        RowMajor, 
        RowMajor, 
        wholeM, 
        wholeK, 
        wholeM, 
        batchCount, 
        path_A2, 
        path_V, 
        path_C2, 
        path_alpha, 
        path_beta, 
        path_validM, 
        path_expect_result2, 
        pre_output2, 
        cur_output2, 
        2, 
        0
    );

    // =========================================================================

    std::string path_A3 = "";
    std::string path_Parameters = src + "Parameters.bin";
    std::string path_C3 = src + "C3.bin";
    std::string path_expect_result3 = src + "expect_result3.bin";

    half* pre_output3 = cur_output2;
    half* cur_output3 = nullptr;

    run_main(
        LLMsGEMM_batch_QKTVP_host,
        stream, 
        RowMajor, 
        RowMajor, 
        wholeM, 
        wholeN, 
        wholeK, 
        batchCount, 
        path_A3, 
        path_Parameters, 
        path_C3, 
        path_alpha, 
        path_beta, 
        path_validM, 
        path_expect_result3, 
        pre_output3, 
        cur_output3, 
        3, 
        0
    );

    // =========================================================================

    ACL_CHECK(aclrtFree(cur_output3));

    // 反初始化
    ACL_CHECK(aclrtDestroyStream(stream));
    ACL_CHECK(aclrtResetDevice(device));
    ACL_CHECK(aclFinalize());
    

    return 0;

}